import paddle
from paddle.vision import transforms as T
import pickle
from build_vocab import Vocabulary

def preprocess_fn(img):
    """
    Takes in an image, Resize it to (224, 224), Normalize if, and Transpose
    it from HWC (shape (224, 224, 3)) to CHW (shape (3, 224, 224)), and
    finally returns the image tensor.
    """
    transform = T.Compose([
        T.Resize(size=(224, 224)), #把数据长宽像素调成224*224
        T.Normalize(mean=[127.5, 127.5, 127.5], std=[127.5, 127.5, 127.5], data_format='HWC'), #标准化
        T.Transpose(), #原始数据形状维度是HWC格式，经过Transpose，转换为CHW格式
        ])
    img = transform(img)
    return paddle.to_tensor(img, dtype='float32')


def collate_fn(data: list):
    # Sort the list of data by caption in descending order
    data.sort(key=lambda x: len(x[1]), reverse=True)
    images, captions = zip(*data)

    # Merge images from a tuple of 3D tensor to a 4D tensor
    images = paddle.stack(images, axis=0)

    # Merge captions from a tuple of 1D tensors with ragged lengths
    # to a 2D tensor
    lengths = [len(cap) for cap in captions]
    targets = paddle.zeros(shape=(len(captions), max(lengths)),
            dtype='int64')
    for i, cap in enumerate(captions):
        end = lengths[i]
        targets[i, :end] = cap[:end]
    
    return images, targets, lengths

def get_vocab(vocab_path):
    with open(vocab_path, 'rb') as f:
        vocab = pickle.load(f)
    return vocab
